//
// ZCInterceptionManager.m
// Copyright 2010 Mark Buer
//
// This file is part of zc-aspect.
//
// zc-aspect is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// zc-aspect is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with zc-aspect.  If not, see <http://www.gnu.org/licenses/>.
//


#import <objc/objc-class.h>
#import "ZCAspect.h"


/*
 TODO
 Make thunkCount be a constant defined in a header file
 Declare structs prior to defining them
 Fix imports (esp uikit)
 Fix standard comments at top of files
 Comment code
 Exception handling -- ie What to do when we cannot intercept?
 Logging?
 */

@implementation ZCInterceptionManager



#define MAX_DESCRIPTORS 1000
static InterceptionDescriptor interceptionDescriptors[MAX_DESCRIPTORS];


+ (void)initialize {
}


+ (int)findFreeDescriptorIndex {
	for (int i = 0; i < MAX_DESCRIPTORS; i++) {
		if (!interceptionDescriptors[i].previousImplementation) {
			return i;
		}
	}
	return -1;
}


extern int thunkSize;
void thunkPool();


+ (void)watchClass:(Class)theClass withInterceptor:(id <ZCInterceptor>)interceptor excludingSelectors:(NSArray *)selectorsToExclude {
	unsigned int methodCount;
	
	Method *methods = class_copyMethodList(theClass, &methodCount);
	
	for (unsigned int i = 0; i < methodCount; i++) {
		Method method = methods[i];
		
		SEL selector = method_getName(method);
		NSNumber *k = [[NSNumber alloc] initWithInt:(unsigned int)selector];
		if (![selectorsToExclude containsObject:k]) {
			[self interceptMethod:method withInterceptor:interceptor];
		}
		[k release];
	}
	
	free(methods);
}


+ (void)interceptMethod:(Method)method withInterceptor:(id <ZCInterceptor>)interceptor {
	const char *typeEncoding = method_getTypeEncoding(method);
	int index = [self findFreeDescriptorIndex];
	if (index != -1) {
		InterceptionDescriptor *descriptor = &interceptionDescriptors[index];

		NSLog(@"Setting up intercept for method: %s", sel_getName(method_getName(method)));
		NSLog(@"-type: %s", typeEncoding);

		descriptor->previousImplementation = method_getImplementation(method);
		descriptor->interceptor = [interceptor retain];
		
		NSMethodSignature *methodSignature = [NSMethodSignature signatureWithObjCTypes:method_getTypeEncoding(method)];
		descriptor->frameSize = [methodSignature frameLength];
		
		char returnType = *[methodSignature methodReturnType];
		switch (returnType) {
			case 'v':
				descriptor->roughReturnType = RRTVoid;
				break;
			case 'c':	// char
			case 'i':	// int
			case 's':	// short
			case 'l':	// long
			case 'C':	// unsigned char
			case 'I':	// unsigned int
			case 'S':	// unsigned short
			case 'L':	// unsigned long
			case 'B':	// _Bool (C99)
			case '*':	//  char *
			case '@':	// id
			case '#':	// Class
			case ':':	// SEL
//check			case '[':	// array
//check			case 'b':	// bitfield
			case '^':	// pointer
//check			case '?':	// unknown (includes function pointers)
				descriptor->roughReturnType = RRTScaler32;
				break;
			case 'q':	// long long
			case 'Q':	// unsigned long long
				descriptor->roughReturnType = RRTScaler64;
				break;
			case 'f':	// float
				descriptor->roughReturnType = RRTFloat32;
				break;
			case 'd':	// double
				descriptor->roughReturnType = RRTFloat64;
				break;
			case 'D':	// long double
				descriptor->roughReturnType = RRTFloat128;
				break;
			case '{':	// struct
			case '(':	// union
				descriptor->roughReturnType = RRTStruct;
				break;
			default:
				descriptor->previousImplementation = nil;
				NSLog(@"Cannot intercept method: %s, unhandled return type", sel_getName(method_getName(method)));
				NSLog(@"-type: %s", typeEncoding);					
				break;
		}
		
		IMP newImplementation = (IMP)(&thunkPool + (index * thunkSize));
		method_setImplementation(method, newImplementation);
	} else {
		NSLog(@"Cannot intercept method: %s, no thunks remaining in pool", sel_getName(method_getName(method)));
		NSLog(@"-type: %s", typeEncoding);
	}

}


void commonInterceptor() {
	register int index __asm__("eax");	// eax is a "disposable" register
	InterceptionDescriptor *interceptionDescriptor = &interceptionDescriptors[index];
	NSLog(@"Intercepted using thunk number %i", index);
	
	Invocation invocation;
	invocation.interceptionDescriptor = interceptionDescriptor;
	invocation.stackFrameBase = __builtin_frame_address(0);
	
	//
	// Call out to the "mid-level" interceptor
	//
	[interceptionDescriptor->interceptor interceptInvocation:&invocation];
	
	//
	// Restore return value (if not struct or void) as stored in the invocation
	//
	register void *returnValue __asm__("esi") = &invocation.returnValue;
	switch (interceptionDescriptor->roughReturnType) {
		case RRTScaler32:
		{
			__asm__("movl (%[source]), %%eax"
					:
					: [source] "r" (returnValue));			
			break;
		}
		case RRTScaler64:
		{
			__asm__("movl (%[source]), %%eax\n"
					"movl 0x4(%[source]), %%edx"
					:
					: [source] "r" (returnValue));			
			break;
		}
		case RRTFloat32:
		{
			__asm__("flds (%[source])"
					:
					: [source] "r" (returnValue));			
			break;
		}
		case RRTFloat64:
		{
			__asm__("fldl (%[source])"
					:
					: [source] "r" (returnValue));			
			break;
		}
		case RRTFloat128:
		{
			__asm__("fldt (%[source])"
					:
					: [source] "r" (returnValue));			
			break;
		}
	}	
}

#define REMOVE_LOW_16_MASK (~0xf)

void invocation_invoke(Invocation *invocation) {
	InterceptionDescriptor *interceptionDescriptor = invocation->interceptionDescriptor;
	
	//
	// Save stack pointer
	//
	void *savedStackPointer;
	__asm__("movl %%esp, %[destination]"
			: [destination] "=rm" (savedStackPointer));
	
	//
	// Adjust stack pointer:
	// * Reserve enough space for the original invocation arguments
	// * Ensure 16 byte alignment so as to be ABI compliant
	//	
	register void *stackPointerAdjustment;
	stackPointerAdjustment = (void *)(REMOVE_LOW_16_MASK & (interceptionDescriptor->frameSize + 0xf));
	__asm__("subl %[source], %%esp"
			:
			: [source] "r" (stackPointerAdjustment));
	
	//
	// Copy invocation arguments from original stack frame to new stack frame
	//
	__asm__(
			"shrl $2, %[length]\n"
			"addl $8, %[sourceAddress]\n"
			"movl %%esp, %%edi\n"
			"cld\n"
			"rep movsd"
			:
			: [sourceAddress] "S" (invocation->stackFrameBase),
			  [length] "c" (interceptionDescriptor->frameSize)
			: "edi"
			);

	//
	// Call the original implementation and save the return value
	//
	void *previousImplementation = interceptionDescriptor->previousImplementation;
	register void *returnValue __asm__("edi") = &invocation->returnValue;
	switch (interceptionDescriptor->roughReturnType) {
		case RRTVoid:
		{
			((void (*)())previousImplementation)();
			break;
		}
		case RRTStruct:
		{
			((void (*)())previousImplementation)();
			break;
		}
		case RRTScaler32:
		{
			*((uint32_t *)returnValue) = ((uint32_t (*)())previousImplementation)();
			break;
		}
		case RRTScaler64:
		{
			*((uint64_t *)returnValue) = ((uint64_t (*)())previousImplementation)();
			break;
		}
		case RRTFloat32:
		{
			*((float *)returnValue) = ((float (*)())previousImplementation)();
			break;
		}
		case RRTFloat64:
		{
			*((double *)returnValue) = ((double (*)())previousImplementation)();
			break;
		}
		case RRTFloat128:
		{
			*((long double *)returnValue) = ((long double (*)())previousImplementation)();
			break;
		}
	}

	//
	// Restore stack pointer
	//
	__asm__("movl %[source], %%esp"
			:
			: [source] "rm" (savedStackPointer));
}


@end
